# 导入必要的库
import torch
import torchvision.transforms as transforms
from PIL import Image
from load_resnet import load_pretrained_resnet


def load_imagenet_labels(filename="imagenet_classes.txt"):
    """从给定文件名读取ImageNet的标签"""
    with open(filename, 'r') as f:
        labels = [line.strip() for line in f.readlines()]
    return labels


def serve_chicken(image_path):
    """
    接收一个图像路径，用ResNet模型进行预测，然后返回预测结果。

    参数:
    - image_path: 输入图像的路径

    返回:
    - prediction: 预测的类别名
    """
    # 确定使用的设备: GPU (如果可用) 或 CPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 先准备烤鸡
    model = load_pretrained_resnet()

    # 为图像制作"餐具" (transforms)
    preprocess = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    # 打开图像文件并准备"上菜"
    image = Image.open(image_path).convert("RGB")
    image = preprocess(image).unsqueeze(0).to(device)  # 确保图像在同一设备上

    # 使用烤鸡为客人"上菜" (预测图像)
    with torch.no_grad():
        outputs = model(image)
        _, predicted = outputs.max(1)
        # 假设你有一个标签到名字的映射，这里只是一个示例
        labels = load_imagenet_labels()
        prediction = labels[predicted[0]]

    return prediction


# 如果这个脚本被当作主程序运行，那么进行预测并打印结果
if __name__ == "__main__":
    result = serve_chicken("headimgurl.jpg")
    print(f"这只烤鸡味道像...{result}!")